import numpy as np;  # @UnusedImport
import pandas as pd;
import matplotlib.pyplot as pl;
import cv2;

def set_Tracks2Dto3D():
    sDir = "C:/Users/Alfredo/workspace/Visual_B3/src/StoredData/Tracks/";
    sTracksFile = sDir + "Tracks_B3.csv";
    dfTracks = pd.read_csv(sTracksFile);
    print dfTracks;
    for i in dfTracks.index:
        aPos = dfTracks.values[i, 2:4]; #Get Columns PosX, PosY
        aPos *= 10;
        dfTracks["PosX"][i] = aPos[0];
        dfTracks["PosY"][i] = aPos[1];
    dfTracks.to_csv(sDir + "Test_3Dto2D.csv", index=False);
    print dfTracks;
    
def get2Dto3D(aP, aR, aT, nF):
    def c(n): return np.cos(n);
    def s(n): return np.sin(n);
    nSx, nSy, nCx, nCy = 1,1,0,0;
    nRx, nRy, nRz = aR;
    nTx, nTy, nTz = aT;
    mInt = np.array([[-nF/nSx, 0, nCx], [0, -nF/nSy, nCy], [0,0,1]], np.float);
    mExt = np.array([[c(nRx)*c(nRy), c(nRx)*s(nRy)*s(nRz)-c(nRz)*s(nRx), 
                      s(nRx)*s(nRz)+c(nRx)*c(nRz)*s(nRy)],
                     [c(nRy)*s(nRx), c(nRx)*c(nRz)+s(nRx)*s(nRy)*s(nRz),
                      c(nRz)*s(nRx)*s(nRy)-c(nRx)*s(nRz)],
                     [-s(nRy), c(nRy)*s(nRz), c(nRy)*c(nRz)]], np.float);
    mT = np.array([[nTx, nTy, nTz]], np.float).T;
    mExt = np.hstack([mExt, mT]);
    aP = np.vstack([aP, np.ones((1,np.shape(aP)[1]))]);
    aP_out = np.dot(np.dot(mInt, mExt), aP);
    aPz = np.vstack([aP_out[2,:],aP_out[2,:]]).T;
    aP_out = np.multiply(aP_out.T[:,:2],aPz);
    return aP_out;
    
    
def plot_Cube():
    #Build Square
    aPts3D = np.zeros((8,3));
    aPts3D[0,:] = np.array([0,0,0]);
    aPts3D[1,:] = np.array([0,1,0]);
    aPts3D[2,:] = np.array([1,1,0]);
    aPts3D[3,:] = np.array([1,0,0]);
    aPts3D[4,:] = np.array([0,0,1]);
    aPts3D[5,:] = np.array([0,1,1]);
    aPts3D[6,:] = np.array([1,1,1]);
    aPts3D[7,:] = np.array([1,0,1]);
    aEdges = [];
    aEdges.append([0,1]);
    aEdges.append([1,2]);
    aEdges.append([2,3]);
    aEdges.append([3,0]);
    aEdges.append([4,5]);
    aEdges.append([5,6]);
    aEdges.append([6,7]);
    aEdges.append([7,4]);
    aEdges.append([0,4]);
    aEdges.append([1,5]);
    aEdges.append([2,6]);
    aEdges.append([3,7]);
    #Transform to 2D
#     pi = np.pi;
#     aR = np.array([0,0,0], np.float); #Rotation
#     aT = np.array([0,0,-2], np.float); #Translation
#     nFx, nFy, nCx, nCy = 1, 1, 0, 0;
#     aCam = np.array([[nFoV[0],0,nCx],[0,nFoV[0],nCy],[0,0,1]], np.float);
    aPts2D = get2Dto3D(aPts3D.T, aR, aT, nFoV[0]);  # @UnusedVariable
#     (aPts2D,NA) = cv2.projectPoints(aPts3D, aR, aT, aCam, None);  # @UnusedVariable
#     aPts2D = np.reshape(aPts2D, (8,2));
#     print aPts3D;
#     print aPts2D;
    #Plot 2D-Square
    pl.hold(True);
    for e in aEdges:
        p0, p1 = e;
        aX = [aPts2D[p0,0], aPts2D[p1,0]];
        aY = [aPts2D[p0,1], aPts2D[p1,1]];
        pl.plot(aX, aY, 'bo-');    
    for i,p in enumerate(aPts2D):
        pl.plot(p[0], p[1], 'r');
        pl.gca().annotate(str(i)+'.'+str(tuple(aPts3D[i,:].astype(np.int))), 
                          xy=tuple(p), xytext=tuple(p+0.01));
    nP = 10;
    pl.axis([-nP, nP, -nP, nP]);
#     pl.show();

def plot_Square_3D():
    #Build Square
    aPts3D = np.zeros((8,3));
    aPts3D[0,:] = np.array([-1,0,1]);
    aPts3D[1,:] = np.array([1,0,1]);
    aPts3D[2,:] = np.array([1,0,-1]);
    aPts3D[3,:] = np.array([-1,0,-1]);
    #Build Axis
    aPts3D[4,:] = np.array([0,0,0]);
    aPts3D[5,:] = np.array([1,0,0]);
    aPts3D[6,:] = np.array([0,1,0]);
    aPts3D[7,:] = np.array([0,0,1]);
    aEdges = [];
    aEdges.append([0,1]);
    aEdges.append([1,2]);
    aEdges.append([2,3]);
    aEdges.append([3,0]);
    #Axis    
    aEdges.append([4,5]);
    aEdges.append([4,6]);
    aEdges.append([4,7]);
    #Transform to 2D
#     pi = np.pi;
#     aR = np.array([-pi*0.5,0,0], np.float); #Rotation
#     aT = np.array([0,-2,0], np.float); #Translation
#     nFx, nFy, nCx, nCy = 1, 1, 0, 0;
#     aCam = np.array([[nFoV[0],0,nCx],[0,nFoV[0],nCy],[0,0,1]], np.float);
    aPts2D = get2Dto3D(aPts3D.T, aR, aT, nFoV[0]);  # @UnusedVariable
#     (aPts2D,NA) = cv2.projectPoints(aPts3D, aR, aT, aCam, None);  # @UnusedVariable
#     aPts2D = np.reshape(aPts2D, (8,2));
#     print aPts3D;
#     print aPts2D;
    #Plot 2D-Square
    pl.hold(True);
    for i,e in enumerate(aEdges):
        p0, p1 = e;
        aX = [aPts2D[p0,0], aPts2D[p1,0]];
        aY = [aPts2D[p0,1], aPts2D[p1,1]];
        if i<4:
            ax.plot(aX, aY, 'bo-');
        else:
            ax.plot(aX, aY, 'r*-');
    sA = ['O', 'X','Y','Z'];
    for i,p in enumerate(aPts2D):
        ax.plot(p[0], p[1], 'r');
        if i<4:
            ax.annotate(str(i)+'.'+str(tuple(aPts3D[i,:].astype(np.int))), 
                              xy=tuple(p), xytext=tuple(p+0.01));
        else:
#             print i-4;
            ax.annotate(sA[i-4], xy=tuple(p), xytext=tuple(p+0.01));
    nP = 20;
#     ax.axis('tight');
    ax.axis([-nP, nP, -nP, nP]);
#     pl.show();
    
def press(event):    
#     print'press', event.key;
    nInc = 0.1;
    if event.key=='up': aR[0]+=nInc;
    elif event.key=='down': aR[0]-=nInc;
    elif event.key=='left': aR[1]+=nInc;
    elif event.key=='right': aR[1]-=nInc;
    elif event.key=='pageup': aR[2]+=nInc;
    elif event.key=='pagedown': aR[2]-=nInc;
    elif event.key=='5': aT[1]+=nInc;
    elif event.key=='2': aT[1]-=nInc;
    elif event.key=='1': aT[0]+=nInc;
    elif event.key=='3': aT[0]-=nInc;
    elif event.key=='7': aT[2]+=nInc;
    elif event.key=='4': aT[2]-=nInc;
    elif event.key=='n': nFoV[0]+=nInc;
    elif event.key=='m': nFoV[0]-=nInc; 
    elif event.key=='+': nInc+=0.1;
    elif event.key=='-': nInc-=0.1;
    ax.cla();
#     plot_Cube();
    plot_Square_3D();
    ax.figure.canvas.draw();
    print list(aR), list(aT);
    
if __name__ == '__main__':
    aR = np.array([0,0,0], np.float); #Rotation
    aT = np.array([0,0,0], np.float); #Translation
    nFoV = np.array([2], np.float);
    nCx, nCy = 0., 0.;
    fig = pl.figure();
    ax = fig.add_subplot(111);
    ax.figure.canvas.mpl_connect('key_press_event', press);
#     set_Tracks2Dto3D();
#     plot_Square_2D();
    #pl.hold(False);
    plot_Square_3D();
#     plot_Cube();
    pl.show();
